{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ca2a24f-ddd0-46d6-843d-03f7e7741ac5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dendron import *\n",
    "from dendron.conditions import CompletionConditionConfig, CompletionCondition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "026f9381-c502-404d-a1cd-73a1c618aaf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg = CompletionConditionConfig(load_in_4bit=True,\n",
    "                                    model_name='v1olet/v1olet_merged_dpo_7B',\n",
    "                                    use_flash_attn_2=True)\n",
    "node = CompletionCondition('cond-node', cfg)\n",
    "\n",
    "tree = BehaviorTree(\"completion-condition-tree\", node)\n",
    "tree.pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcaddafe-2532-4773-8cd8-beb84621fbd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"\"\"{system}\n",
    "### Instruction:\n",
    "{query}\n",
    "\n",
    "### Response:\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c89a99c3-a387-4731-9e54-c3ea0edca5ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "def success(completion):\n",
    "    if completion == \"yes\":\n",
    "        return NodeStatus.SUCCESS\n",
    "    else:\n",
    "        return NodeStatus.FAILURE\n",
    "\n",
    "tree.blackboard[\"completions_in\"] = ['yes', 'no']\n",
    "tree.blackboard[\"success_fn\"] = success \n",
    "tree.blackboard[\"in\"] = prompt.format(system=\"\", \n",
    "                                      query=\"The following is a headline from the Wall Street Journal: 'Nation’s Top Economists Are Short-Term Happy, Long-Term Glum.' Is this headline about the economy?\")\n",
    "\n",
    "status = tree.tick_once()\n",
    "status"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25f3fd28-502d-41e0-8aac-bc57a7de30f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def success1(completion):\n",
    "    if completion == \"cats\":\n",
    "        return NodeStatus.SUCCESS\n",
    "    else:\n",
    "        return NodeStatus.FAILURE\n",
    "\n",
    "tree.blackboard[\"completions_in\"] = ['dogs', 'cats']\n",
    "tree.blackboard[\"success_fn\"] = success1\n",
    "tree.blackboard[\"in\"] = prompt.format(system=\"\", query=\"Which animal is more likely to purr: dogs or cats?\")\n",
    "\n",
    "status = tree.tick_once()\n",
    "status"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67d9df8b-e79f-4067-bb0b-484a93d83a8f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
